#include "vio/tools/vision_utils.h"

namespace vio {
namespace vision_utils {

float ShiTomasiScore(const cv::Mat& img, int u, int v) {
  assert(img.type() == CV_8UC1);

  float dXX = 0.0;
  float dYY = 0.0;
  float dXY = 0.0;
  const int halfbox_size = 4;
  const int box_size = 2 * halfbox_size;
  const int box_area = box_size * box_size;
  const int x_min = u - halfbox_size;
  const int x_max = u + halfbox_size;
  const int y_min = v - halfbox_size;
  const int y_max = v + halfbox_size;

  if (x_min < 1 || x_max >= img.cols - 1 || y_min < 1 || y_max >= img.rows - 1)
    return 0.0;  // patch is too close to the boundary

  const int stride = img.step.p[0];
  for (int y = y_min; y < y_max; ++y) {
    const uint8_t* ptr_left = img.data + stride * y + x_min - 1;
    const uint8_t* ptr_right = img.data + stride * y + x_min + 1;
    const uint8_t* ptr_top = img.data + stride * (y - 1) + x_min;
    const uint8_t* ptr_bottom = img.data + stride * (y + 1) + x_min;
    for (int x = 0; x < box_size;
         ++x, ++ptr_left, ++ptr_right, ++ptr_top, ++ptr_bottom) {
      float dx = *ptr_right - *ptr_left;
      float dy = *ptr_bottom - *ptr_top;
      dXX += dx * dx;
      dYY += dy * dy;
      dXY += dx * dy;
    }
  }

  // Find and return smaller eigenvalue:
  dXX = dXX / (2.0 * box_area);
  dYY = dYY / (2.0 * box_area);
  dXY = dXY / (2.0 * box_area);
  return 0.5 * (dXX + dYY -
                sqrt((dXX + dYY) * (dXX + dYY) - 4 * (dXX * dYY - dXY * dXY)));
}

}  // namespace vision_utils
}  // namespace vio
